import math
import operator
from collections import Counter

def euclidean_distance(point1, point2):
    """
    计算两个点之间的欧几里得距离
    """
    if len(point1) != len(point2):
        raise ValueError("Points must have the same dimensions")
    
    squared_distance = 0
    for i in range(len(point1)):
        squared_distance += (point1[i] - point2[i]) ** 2
    
    return math.sqrt(squared_distance)

def manhattan_distance(point1, point2):
    """
    计算两个点之间的曼哈顿距离
    """
    if len(point1) != len(point2):
        raise ValueError("Points must have the same dimensions")
    
    distance = 0
    for i in range(len(point1)):
        distance += abs(point1[i] - point2[i])
    
    return distance

def get_neighbors(training_set, test_instance, k, distance_metric='euclidean'):
    """
    获取测试实例的k个最近邻
    """
    distances = []
    
    # 选择距离度量函数
    if distance_metric == 'euclidean':
        distance_func = euclidean_distance
    elif distance_metric == 'manhattan':
        distance_func = manhattan_distance
    else:
        raise ValueError("Unsupported distance metric")
    
    # 计算测试实例与所有训练实例的距离
    for i, training_point in enumerate(training_set):
        # training_point格式: (features, label)
        features = training_point[0]
        label = training_point[1]
        
        dist = distance_func(test_instance, features)
        distances.append((training_point, dist, label))
    
    # 按距离排序并返回前k个
    distances.sort(key=operator.itemgetter(1))
    neighbors = distances[:k]
    
    return neighbors

def predict_classification(neighbors):
    """
    基于邻居进行分类预测（多数投票）
    """
    class_votes = {}
    
    for neighbor in neighbors:
        label = neighbor[2]  # 邻居的标签
        if label in class_votes:
            class_votes[label] += 1
        else:
            class_votes[label] = 1
    
    # 按票数排序，返回票数最多的类别
    sorted_votes = sorted(class_votes.items(), key=operator.itemgetter(1), reverse=True)
    return sorted_votes[0][0]

def predict_regression(neighbors):
    """
    基于邻居进行回归预测（平均值）
    """
    if not neighbors:
        return 0
    
    # 假设标签是数值型
    values = [neighbor[2] for neighbor in neighbors]
    return sum(values) / len(values)

class KNN:
    """
    K近邻算法类
    """
    
    def __init__(self, k=3, distance_metric='euclidean'):
        """
        初始化KNN模型
        
        参数:
        k: 邻居数量
        distance_metric: 距离度量方法 ('euclidean' 或 'manhattan')
        """
        self.k = k
        self.distance_metric = distance_metric
        self.training_set = None
    
    def fit(self, X_train, y_train):
        """
        训练模型（实际上只是存储训练数据）
        """
        if len(X_train) != len(y_train):
            raise ValueError("X_train and y_train must have the same length")
        
        self.training_set = []
        for i in range(len(X_train)):
            self.training_set.append((X_train[i], y_train[i]))
    
    def predict(self, X_test, problem_type='classification'):
        """
        预测测试数据的标签
        
        参数:
        X_test: 测试数据
        problem_type: 问题类型 ('classification' 或 'regression')
        """
        if self.training_set is None:
            raise ValueError("Model must be fitted before prediction")
        
        predictions = []
        
        for test_instance in X_test:
            neighbors = get_neighbors(self.training_set, test_instance, self.k, self.distance_metric)
            
            if problem_type == 'classification':
                prediction = predict_classification(neighbors)
            elif problem_type == 'regression':
                prediction = predict_regression(neighbors)
            else:
                raise ValueError("problem_type must be 'classification' or 'regression'")
            
            predictions.append(prediction)
        
        return predictions
    
    def score(self, X_test, y_test, problem_type='classification'):
        """
        计算模型在测试集上的准确率（分类）或R²分数（回归）
        """
        predictions = self.predict(X_test, problem_type)
        
        if problem_type == 'classification':
            # 计算准确率
            correct = 0
            for i in range(len(predictions)):
                if predictions[i] == y_test[i]:
                    correct += 1
            return correct / len(predictions)
        
        else:  # regression
            # 计算R²分数
            y_mean = sum(y_test) / len(y_test)
            ss_total = sum((y - y_mean) ** 2 for y in y_test)
            ss_residual = sum((y_test[i] - predictions[i]) ** 2 for i in range(len(y_test)))
            
            if ss_total == 0:
                return 1.0  # 完美预测
            
            return 1 - (ss_residual / ss_total)

# 测试代码
if __name__ == "__main__":
    print("=== K近邻算法测试 ===\n")
    
    # 测试1：分类问题
    print("1. 分类问题测试:")
    
    # 简单的二维分类数据
    X_train_class = [
        [1, 2], [1, 4], [2, 1], [2, 3],  # 类别0
        [5, 6], [6, 5], [7, 7], [6, 8]   # 类别1
    ]
    y_train_class = [0, 0, 0, 0, 1, 1, 1, 1]
    
    X_test_class = [[1, 3], [6, 6]]
    y_test_class = [0, 1]
    
    # 创建并训练KNN分类器
    knn_classifier = KNN(k=3, distance_metric='euclidean')
    knn_classifier.fit(X_train_class, y_train_class)
    
    # 预测
    predictions_class = knn_classifier.predict(X_test_class, 'classification')
    accuracy = knn_classifier.score(X_test_class, y_test_class, 'classification')
    
    print(f"训练数据: {X_train_class}")
    print(f"训练标签: {y_train_class}")
    print(f"测试数据: {X_test_class}")
    print(f"真实标签: {y_test_class}")
    print(f"预测结果: {predictions_class}")
    print(f"准确率: {accuracy:.2f}")
    
    # 测试2：回归问题
    print("\n2. 回归问题测试:")
    
    # 简单的一维回归数据
    X_train_reg = [[1], [2], [3], [4], [5], [6]]
    y_train_reg = [2, 4, 6, 8, 10, 12]  # y = 2x
    
    X_test_reg = [[2.5], [4.5]]
    y_test_reg = [5, 9]
    
    # 创建并训练KNN回归器
    knn_regressor = KNN(k=2, distance_metric='euclidean')
    knn_regressor.fit(X_train_reg, y_train_reg)
    
    # 预测
    predictions_reg = knn_regressor.predict(X_test_reg, 'regression')
    r2_score = knn_regressor.score(X_test_reg, y_test_reg, 'regression')
    
    print(f"训练数据: {X_train_reg}")
    print(f"训练标签: {y_train_reg}")
    print(f"测试数据: {X_test_reg}")
    print(f"真实值: {y_test_reg}")
    print(f"预测值: {[f'{x:.2f}' for x in predictions_reg]}")
    print(f"R²分数: {r2_score:.2f}")
    
    # 测试3：手动使用函数
    print("\n3. 手动函数使用示例:")
    
    # 准备训练数据（格式：[(features, label), ...]）
    training_data = []
    for i in range(len(X_train_class)):
        training_data.append((X_train_class[i], y_train_class[i]))
    
    test_point = [1, 3]
    neighbors = get_neighbors(training_data, test_point, k=3)
    
    print(f"测试点: {test_point}")
    print("最近的3个邻居:")
    for i, (point, distance, label) in enumerate(neighbors):
        print(f"  邻居{i+1}: 点{point[0]}, 距离{distance:.2f}, 标签{label}")
    
    prediction = predict_classification(neighbors)
    print(f"预测标签: {prediction}")